import torch
from torch.utils.data import DataLoader
from data_process import Data_Loader
from model_plus import UNet_ViT
import os
from PIL import Image
import numpy as np

# 确保保存图片的文件夹存在
save_dir = "result"
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

# 加载测试集
testset = Data_Loader("processed_data/test/images")
test_loader = DataLoader(dataset=testset, batch_size=1)

# 加载模型
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
model = UNet_ViT(in_channels=1, num_classes=1, image_size=584, patch_size=128, vit_embed_dim=768, vit_depth=12, vit_heads=12)
model.load_state_dict(torch.load('Best_model/UNet_ViT_128.pth', map_location=device))
model.to(device)
model.eval()

# 初始化指标累加变量
total_ACC = 0
total_PPV = 0
total_TPR = 0
total_F1 = 0
total_JS = 0

# 遍历测试集
with torch.no_grad():
    for i, (image, label) in enumerate(test_loader):
        image, label = image.to(device), label.to(device)
        outputs = model(image)
        outputs = torch.sigmoid(outputs)  # 应用sigmoid激活函数
        outputs = (outputs > 0.5).float()  # 二值化处理

        # 统计TP、FP、FN、TN
        TP_img = ((outputs == 1) & (label == 1)).sum().item()
        FP_img = ((outputs == 1) & (label == 0)).sum().item()
        FN_img = ((outputs == 0) & (label == 1)).sum().item()
        TN_img = ((outputs == 0) & (label == 0)).sum().item()

        # 计算当前图片的各项指标
        ACC_img = (TP_img + TN_img) / (TP_img + FP_img + FN_img + TN_img)
        PPV_img = TP_img / (TP_img + FP_img) if (TP_img + FP_img) != 0 else 0
        TPR_img = TP_img / (TP_img + FN_img) if (TP_img + FN_img) != 0 else 0
        F1_img = 2 * TP_img / (2 * TP_img + FP_img + FN_img) if (2 * TP_img + FP_img + FN_img) != 0 else 0
        JS_img = TP_img / (TP_img + FP_img + FN_img) if (TP_img + FP_img + FN_img) != 0 else 0

        # 累加指标值
        total_ACC += ACC_img
        total_PPV += PPV_img
        total_TPR += TPR_img
        total_F1 += F1_img
        total_JS += JS_img

        # 打印当前图片的结果
        print(f"Image {i+1} - ACC: {ACC_img:.4f}, PPV: {PPV_img:.4f}, TPR: {TPR_img:.4f}, F1: {F1_img:.4f}, JS: {JS_img:.4f}")

# 计算平均指标
avg_ACC = total_ACC / len(test_loader)
avg_PPV = total_PPV / len(test_loader)
avg_TPR = total_TPR / len(test_loader)
avg_F1 = total_F1 / len(test_loader)
avg_JS = total_JS / len(test_loader)

# 打印平均结果
print(f"Average - ACC: {avg_ACC:.4f}, PPV: {avg_PPV:.4f}, TPR: {avg_TPR:.4f}, F1: {avg_F1:.4f}, JS: {avg_JS:.4f}")